#ifndef WRAPPER_H_
#define WRAPPER_H_

#include <array>
#include <memory>
#include <mutex>
#include <string>
#include <utility>

#include "carla/client/ActorList.h"
#include "carla/client/BlueprintLibrary.h"
#include "carla/client/Client.h"
#include "carla/client/Map.h"
#include "carla/client/Sensor.h"
#include "carla/client/Vehicle.h"
#include "carla/sensor/data/CollisionEvent.h"
#include "carla/sensor/data/GnssMeasurement.h"
#include "carla/sensor/data/Image.h"
#include "carla/sensor/data/LaneInvasionEvent.h"
#include "rust/cxx.h"

struct Control {
  float get_throttle() const { return throttle; }

  float get_steer() const { return steer; }

  float get_brake() const { return brake; }

  bool get_hand_brake() const { return hand_brake; }

  bool get_reverse() const { return reverse; }

  bool get_manual_gear_shift() const { return manual_gear_shift; }

  int32_t get_gear() const { return gear; }

  float throttle = 0.0f;
  float steer = 0.0f;
  float brake = 0.0f;
  bool hand_brake = false;
  bool reverse = false;
  bool manual_gear_shift = false;
  int32_t gear = 0;
};

class ActorSnapshot {
 public:
  explicit ActorSnapshot(carla::client::ActorSnapshot&& snapshot)
      : inner_(std::move(snapshot)) {}

  explicit ActorSnapshot(const carla::client::ActorSnapshot& snapshot)
      : inner_(snapshot) {}

  uint32_t get_id() const { return static_cast<uint32_t>(inner_.id); }

  std::array<float, 3> get_location() const {
    const auto& location = inner_.transform.location;
    return {location.x, location.y, location.z};
  }

  std::array<float, 3> get_rotation() const {
    const auto& rotation = inner_.transform.rotation;
    return {rotation.pitch, rotation.yaw, rotation.roll};
  }

  std::array<float, 3> get_velocity() const {
    const auto& v = inner_.velocity;
    return {v.x, v.y, v.z};
  }

  std::array<float, 3> get_angular_velocity() const {
    const auto& v = inner_.velocity;
    return {v.x, v.y, v.z};
  }

  std::array<float, 3> get_acceleration() const {
    const auto& v = inner_.acceleration;
    return {v.x, v.y, v.z};
  }

  std::unique_ptr<Control> get_control() const {
    auto control =
        carla::client::Vehicle::Control(inner_.state.vehicle_data.control);

    Control r;
    r.throttle = control.throttle;
    r.steer = control.steer;
    r.brake = control.brake;
    r.hand_brake = control.hand_brake;
    r.reverse = control.reverse;
    r.manual_gear_shift = control.manual_gear_shift;
    r.gear = control.gear;

    return std::make_unique<Control>(std::move(r));
  }

  float get_speed_limit() const {
    return inner_.state.vehicle_data.speed_limit;
  }

  bool has_traffic_light() const {
    return inner_.state.vehicle_data.has_traffic_light;
  }

  uint32_t get_traffic_light_id() const {
    return static_cast<uint32_t>(inner_.state.vehicle_data.traffic_light_id);
  }

  int32_t get_traffic_light_state() const {
    return static_cast<int32_t>(inner_.state.vehicle_data.traffic_light_state);
  }

 private:
  carla::client::ActorSnapshot inner_;
};

class Actor {
 public:
  explicit Actor(carla::SharedPtr<carla::client::Actor>&& actor)
      : inner_(std::move(actor)) {}

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  rust::String get_type_id() const { return inner_->GetTypeId(); }

  bool destroy() { return inner_->Destroy(); }

  rust::Vec<uint8_t> get_semantic_tags() const {
    auto tags = inner_->GetSemanticTags();

    rust::Vec<uint8_t> r;
    r.reserve(tags.size());
    for (auto tag : tags) {
      r.push_back(tag);
    }

    return r;
  }

  std::array<float, 9> get_bounding_box() const {
    const auto& bbox = inner_->GetBoundingBox();
    return {bbox.location.x,     bbox.location.y,   bbox.location.z,
            bbox.rotation.pitch, bbox.rotation.yaw, bbox.rotation.roll,
            bbox.extent.x,       bbox.extent.y,     bbox.extent.z};
  }

 private:
  carla::SharedPtr<carla::client::Actor> inner_;
};

class Vehicle {
 public:
  explicit Vehicle(carla::SharedPtr<carla::client::Vehicle>&& vehicle)
      : inner_(std::move(vehicle)) {}

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  rust::String get_type_id() const { return inner_->GetTypeId(); }

  bool destroy() { return inner_->Destroy(); }

  void enable_carsim(const rust::str simfile_path) {
    inner_->EnableCarSim(std::string(simfile_path));
  }

  void set_transform(float x, float y, float z, float pitch, float yaw,
                     float roll) {
    carla::geom::Location location(carla::geom::Vector3D(x, y, z));
    carla::geom::Rotation rotation(pitch, yaw, roll);
    carla::geom::Transform transform(location, rotation);

    inner_->SetTransform(transform);
  }

  void apply_control(float throttle, float steer, float brake, bool hand_brake,
                     bool reverse, bool manual_gear_shift, int32_t gear) {
    carla::client::Vehicle::Control control;
    control.throttle = throttle;
    control.steer = steer;
    control.brake = brake;
    control.hand_brake = hand_brake;
    control.reverse = reverse;
    control.manual_gear_shift = manual_gear_shift;
    control.gear = gear;

    inner_->ApplyControl(control);
  }

 private:
  carla::SharedPtr<carla::client::Vehicle> inner_;
};

struct Localization {
  uint64_t get_timestamp() const {
    return static_cast<unsigned long long>(data->GetTimestamp() * 1e6);
  }

  std::array<double, 3> get_lla() const {
    auto geo =
        boost::static_pointer_cast<carla::sensor::data::GnssMeasurement>(data);
    return {geo->GetLatitude(), geo->GetLongitude(), geo->GetAltitude()};
  }

  std::array<float, 3> get_location() const {
    auto tf =
        boost::static_pointer_cast<carla::sensor::data::GnssMeasurement>(data)
            ->GetSensorTransform();
    return {tf.location.x, tf.location.y, tf.location.z};
  }

  std::array<float, 3> get_rotation() const {
    auto tf =
        boost::static_pointer_cast<carla::sensor::data::GnssMeasurement>(data)
            ->GetSensorTransform();
    return {tf.rotation.pitch, tf.rotation.yaw, tf.rotation.roll};
  }

  carla::SharedPtr<carla::sensor::SensorData> data;
};

class GNSS {
  struct Data {
    std::mutex mtx;
    bool valid = false;
    Localization localization;
  };

 public:
  explicit GNSS(carla::SharedPtr<carla::client::Sensor>&& sensor)
      : inner_(std::move(sensor)) {
    data_.reset(new Data());

    inner_->Listen([data{data_}](auto sensor_data) {
      auto geo =
          boost::static_pointer_cast<carla::sensor::data::GnssMeasurement>(
              sensor_data);

      {
        std::unique_lock<decltype(data->mtx)> lk(data->mtx);

        data->valid = true;
        data->localization.data = std::move(sensor_data);
      }
    });
  }

  ~GNSS() { inner_->Stop(); }

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  std::unique_ptr<Localization> get_localization() {
    std::unique_lock<decltype(data_->mtx)> lk(data_->mtx);

    if (data_->valid) {
      data_->valid = false;
      return std::make_unique<Localization>(std::move(data_->localization));
    } else {
      return nullptr;
    }
  }

 private:
  carla::SharedPtr<carla::client::Sensor> inner_;
  std::shared_ptr<Data> data_;
};

struct Image {
  uint64_t get_timestamp() const {
    return static_cast<unsigned long long>(data->GetTimestamp() * 1e6);
  }

  uint32_t get_width() const {
    auto image = boost::static_pointer_cast<carla::sensor::data::Image>(data);
    return image->GetWidth();
  }

  uint32_t get_height() const {
    auto image = boost::static_pointer_cast<carla::sensor::data::Image>(data);
    return image->GetHeight();
  }

  float get_fov_angle() const {
    auto image = boost::static_pointer_cast<carla::sensor::data::Image>(data);
    return image->GetFOVAngle();
  }

  rust::Slice<const uint8_t> get_bgra_pixels() const {
    auto image = boost::static_pointer_cast<carla::sensor::data::Image>(data);
    return rust::Slice<const uint8_t>(reinterpret_cast<uint8_t*>(image->data()),
                                      image->size() * 4);
  }

  carla::SharedPtr<carla::sensor::SensorData> data;
};

class Camera {
  struct Data {
    std::mutex mtx;
    bool valid = false;
    Image image;
  };

 public:
  explicit Camera(carla::SharedPtr<carla::client::Sensor>&& sensor)
      : inner_(std::move(sensor)) {
    data_.reset(new Data());

    inner_->Listen([data{data_}](auto sensor_data) {
      {
        std::unique_lock<decltype(data->mtx)> lk(data->mtx);

        data->valid = true;
        data->image.data = std::move(sensor_data);
      }
    });
  }

  ~Camera() { inner_->Stop(); }

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  std::unique_ptr<Image> get_image() {
    std::unique_lock<decltype(data_->mtx)> lk(data_->mtx);

    if (data_->valid) {
      data_->valid = false;
      return std::make_unique<Image>(std::move(data_->image));
    } else {
      return nullptr;
    }
  }

 private:
  carla::SharedPtr<carla::client::Sensor> inner_;
  std::shared_ptr<Data> data_;
};

struct CollisionEvent {
  uint64_t get_timestamp() const {
    return static_cast<unsigned long long>(data->GetTimestamp() * 1e6);
  }

  std::unique_ptr<Actor> get_other_actor() const {
    auto event =
        boost::static_pointer_cast<carla::sensor::data::CollisionEvent>(data);
    return std::make_unique<Actor>(event->GetOtherActor());
  }

  std::array<float, 3> get_normal_impulse() const {
    auto v =
        boost::static_pointer_cast<carla::sensor::data::CollisionEvent>(data)
            ->GetNormalImpulse();
    return {v.x, v.y, v.z};
  }

  carla::SharedPtr<carla::sensor::SensorData> data;
};

class CollisionDetector {
  struct Data {
    std::mutex mtx;
    bool valid = false;
    CollisionEvent event;
  };

 public:
  explicit CollisionDetector(carla::SharedPtr<carla::client::Sensor>&& sensor)
      : inner_(std::move(sensor)) {
    data_.reset(new Data());

    inner_->Listen([data{data_}](auto sensor_data) {
      {
        std::unique_lock<decltype(data->mtx)> lk(data->mtx);

        data->valid = true;
        data->event.data = std::move(sensor_data);
      }
    });
  }

  ~CollisionDetector() { inner_->Stop(); }

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  std::unique_ptr<CollisionEvent> get_event() {
    std::unique_lock<decltype(data_->mtx)> lk(data_->mtx);

    if (data_->valid) {
      data_->valid = false;
      return std::make_unique<CollisionEvent>(std::move(data_->event));
    } else {
      return nullptr;
    }
  }

 private:
  carla::SharedPtr<carla::client::Sensor> inner_;
  std::shared_ptr<Data> data_;
};

struct LaneInvasionEvent {
  uint64_t get_timestamp() const {
    return static_cast<unsigned long long>(data->GetTimestamp() * 1e6);
  }

  rust::Vec<int32_t> get_lane_change() const {
    auto event =
        boost::static_pointer_cast<carla::sensor::data::LaneInvasionEvent>(
            data);
    auto lms = event->GetCrossedLaneMarkings();

    rust::Vec<int32_t> r;
    r.reserve(lms.size());
    for (const auto& lm : lms) {
      r.push_back(static_cast<int32_t>(lm.lane_change));
    }

    return r;
  }

  carla::SharedPtr<carla::sensor::SensorData> data;
};

class LaneInvasionDetector {
  struct Data {
    std::mutex mtx;
    bool valid = false;
    LaneInvasionEvent event;
  };

 public:
  explicit LaneInvasionDetector(
      carla::SharedPtr<carla::client::Sensor>&& sensor)
      : inner_(std::move(sensor)) {
    data_.reset(new Data());

    // inner_->Listen([data{data_}](auto sensor_data) {
    //   {
    //     std::unique_lock<decltype(data->mtx)> lk(data->mtx);

    //     data->valid = true;
    //     data->event.data = std::move(sensor_data);
    //   }
    // });
  }

  ~LaneInvasionDetector() { inner_->Stop(); }

  uint32_t get_id() const { return static_cast<uint32_t>(inner_->GetId()); }

  std::unique_ptr<LaneInvasionEvent> get_event() {
    std::unique_lock<decltype(data_->mtx)> lk(data_->mtx);

    if (data_->valid) {
      data_->valid = false;
      return std::make_unique<LaneInvasionEvent>(std::move(data_->event));
    } else {
      return nullptr;
    }
  }

 private:
  carla::SharedPtr<carla::client::Sensor> inner_;
  std::shared_ptr<Data> data_;
};

class Map {
 public:
  explicit Map(carla::SharedPtr<carla::client::Map>&& map)
      : inner_(std::move(map)) {}

  rust::String get_name() const { return inner_->GetName(); }

 private:
  carla::SharedPtr<carla::client::Map> inner_;
};

class WorldSnapshot {
 public:
  explicit WorldSnapshot(carla::client::WorldSnapshot&& snapshot)
      : inner_(std::move(snapshot)) {}

  uint64_t get_running_time() const {
    return static_cast<uint64_t>(inner_.GetTimestamp().elapsed_seconds * 1e6);
  }

  uint64_t get_system_time() const {
    return static_cast<uint64_t>(inner_.GetTimestamp().platform_timestamp *
                                 1e6);
  }

  std::unique_ptr<ActorSnapshot> find_actor_snapshot(uint32_t id) const {
    auto o = inner_.Find(id);
    if (o) {
      return std::make_unique<ActorSnapshot>(std::move(*o));
    } else {
      return nullptr;
    }
  }

  rust::Vec<uint32_t> get_actor_ids() const {
    rust::Vec<uint32_t> r;
    r.reserve(inner_.size());
    for (const auto& actor_snapshot : inner_) {
      r.push_back(static_cast<uint32_t>(actor_snapshot.id));
    }

    return r;
  }

 private:
  carla::client::WorldSnapshot inner_;
};

class World {
 public:
  explicit World(carla::client::World&& world) : inner_(std::move(world)) {}

  std::unique_ptr<WorldSnapshot> get_snapshot() const {
    return std::make_unique<WorldSnapshot>(inner_.GetSnapshot());
  }

  std::unique_ptr<Map> get_map() const {
    return std::make_unique<Map>(inner_.GetMap());
  }

  std::unique_ptr<Actor> spawn_actor(const rust::Str blueprint, float x,
                                     float y, float z, float pitch, float yaw,
                                     float roll) {
    auto bp = inner_.GetBlueprintLibrary()->at(std::string(blueprint));
    carla::geom::Location location(carla::geom::Vector3D(x, y, z));
    carla::geom::Rotation rotation(pitch, yaw, roll);
    carla::geom::Transform transform(location, rotation);

    return std::make_unique<Actor>(inner_.SpawnActor(bp, transform));
  }

  std::unique_ptr<Actor> get_actor(uint32_t id) const {
    auto actor = inner_.GetActor(id);
    if (actor) {
      return std::make_unique<Actor>(std::move(actor));
    } else {
      return nullptr;
    }
  }

  std::unique_ptr<Vehicle> spawn_vehicle(const rust::Str vehicle_type, float x,
                                         float y, float z, float pitch,
                                         float yaw, float roll) {
    auto bp = inner_.GetBlueprintLibrary()->at(std::string("vehicle.") +
                                               std::string(vehicle_type));
    carla::geom::Location location(carla::geom::Vector3D(x, y, z));
    carla::geom::Rotation rotation(pitch, yaw, roll);
    carla::geom::Transform transform(location, rotation);

    return std::make_unique<Vehicle>(
        boost::static_pointer_cast<carla::client::Vehicle>(
            inner_.SpawnActor(bp, transform)));
  }

  std::unique_ptr<Vehicle> get_vehicle(uint32_t id) const {
    auto actor = inner_.GetActor(id);
    if (actor) {
      return std::make_unique<Vehicle>(
          boost::static_pointer_cast<carla::client::Vehicle>(std::move(actor)));
    } else {
      return nullptr;
    }
  }

  std::unique_ptr<Vehicle> get_vehicle_by_role(
      const rust::Str role_name) const {
    const auto actors = inner_.GetActors()->Filter("vehicle.*");
    for (const auto& actor : *actors) {
      for (const auto& attr : actor->GetAttributes()) {
        if (attr.GetId() == "role_name") {
          if (attr.GetValue() == std::string(role_name)) {
            return std::make_unique<Vehicle>(
                boost::static_pointer_cast<carla::client::Vehicle>(
                    std::move(decltype(actor)(actor))));
          }
        }
      }
    }

    return nullptr;
  }

  std::unique_ptr<GNSS> create_gnss(uint32_t parent_id, float x, float y,
                                    float z, float pitch, float yaw, float roll,
                                    float sensor_tick) {
    auto parent = inner_.GetActor(parent_id);
    if (parent) {
      auto bpl = inner_.GetBlueprintLibrary();
      auto bp = *bpl->Find("sensor.other.gnss");
      bp.SetAttribute("sensor_tick", std::to_string(sensor_tick));

      auto location = carla::geom::Location(x, y, z);
      auto rotation = carla::geom::Rotation(pitch, yaw, roll);
      carla::geom::Transform tf(location, rotation);

      return std::make_unique<GNSS>(
          boost::static_pointer_cast<carla::client::Sensor>(
              inner_.SpawnActor(bp, tf, parent.get())));
    } else {
      return nullptr;
    }
  }

  std::unique_ptr<Camera> create_camera(
      uint32_t parent_id, float x, float y, float z, float pitch, float yaw,
      float roll, float fov, uint32_t image_size_x, uint32_t image_size_y,
      float iso, float gamma, float shutter_speed, float sensor_tick) {
    auto parent = inner_.GetActor(parent_id);
    if (parent) {
      auto bpl = inner_.GetBlueprintLibrary();
      auto bp = *bpl->Find("sensor.camera.rgb");
      bp.SetAttribute("fov", std::to_string(fov));
      bp.SetAttribute("image_size_x", std::to_string(image_size_x));
      bp.SetAttribute("image_size_y", std::to_string(image_size_y));
      bp.SetAttribute("iso", std::to_string(iso));
      bp.SetAttribute("gamma", std::to_string(gamma));
      bp.SetAttribute("shutter_speed", std::to_string(shutter_speed));
      bp.SetAttribute("sensor_tick", std::to_string(sensor_tick));

      auto location = carla::geom::Location(x, y, z);
      auto rotation = carla::geom::Rotation(pitch, yaw, roll);
      carla::geom::Transform tf(location, rotation);

      return std::make_unique<Camera>(
          boost::static_pointer_cast<carla::client::Sensor>(
              inner_.SpawnActor(bp, tf, parent.get())));
    } else {
      return nullptr;
    }
  }

  std::unique_ptr<CollisionDetector> create_collision_detector(
      uint32_t parent_id, float x, float y, float z, float pitch, float yaw,
      float roll) {
    auto parent = inner_.GetActor(parent_id);
    if (parent) {
      auto bpl = inner_.GetBlueprintLibrary();
      auto bp = *bpl->Find("sensor.other.collision");

      auto location = carla::geom::Location(x, y, z);
      auto rotation = carla::geom::Rotation(pitch, yaw, roll);
      carla::geom::Transform tf(location, rotation);

      return std::make_unique<CollisionDetector>(
          boost::static_pointer_cast<carla::client::Sensor>(
              inner_.SpawnActor(bp, tf, parent.get())));
    } else {
      return nullptr;
    }
  }

  std::unique_ptr<LaneInvasionDetector> create_lane_invasion_detector(
      uint32_t parent_id, float x, float y, float z, float pitch, float yaw,
      float roll) {
    auto parent = inner_.GetActor(parent_id);
    if (parent) {
      auto bpl = inner_.GetBlueprintLibrary();
      auto bp = *bpl->Find("sensor.other.lane_invasion");

      auto location = carla::geom::Location(x, y, z);
      auto rotation = carla::geom::Rotation(pitch, yaw, roll);
      carla::geom::Transform tf(location, rotation);

      return std::make_unique<LaneInvasionDetector>(
          boost::static_pointer_cast<carla::client::Sensor>(
              inner_.SpawnActor(bp, tf, parent.get())));
    } else {
      return nullptr;
    }
  }

  bool destroy_actor_by_id(uint32_t id) {
    auto actor = inner_.GetActor(id);
    if (actor) {
      return actor->Destroy();
    }

    return true;
  }

 private:
  carla::client::World inner_;
};

class Client {
 public:
  explicit Client(carla::client::Client&& client) : inner_(std::move(client)) {}

  void set_timeout(uint32_t milliseconds) {
    inner_.SetTimeout(carla::time_duration::milliseconds(milliseconds));
  }

  std::unique_ptr<World> get_world() const {
    return std::make_unique<World>(inner_.GetWorld());
  }

  std::unique_ptr<World> load_world(const rust::Str map_name,
                                    bool reset_settings) {
    return std::make_unique<World>(inner_.LoadWorld(
        std::string(map_name), reset_settings, carla::rpc::MapLayer::All));
  }

  std::unique_ptr<World> generate_open_drive_world(
      const rust::Str opendrive, double vertex_distance, double max_road_length,
      double wall_height, double additional_width, bool smooth_junctions,
      bool enable_mesh_visibility, bool enable_pedestrian_navigation,
      bool reset_settings) {
    carla::rpc::OpendriveGenerationParameters params;
    params.vertex_distance = vertex_distance;
    params.max_road_length = max_road_length;
    params.wall_height = wall_height;
    params.additional_width = additional_width;
    params.smooth_junctions = smooth_junctions;
    params.enable_mesh_visibility = enable_mesh_visibility;
    params.enable_pedestrian_navigation = enable_pedestrian_navigation;

    return std::make_unique<World>(inner_.GenerateOpenDriveWorld(
        std::string(opendrive), params, reset_settings));
  }

  std::unique_ptr<World> reload_world(bool reset_settings) {
    return std::make_unique<World>(inner_.ReloadWorld(reset_settings));
  }

 private:
  carla::client::Client inner_;
};

std::unique_ptr<Client> new_client(const rust::Str host, uint16_t port,
                                   uint32_t worker_threads);

#endif
